﻿#!/usr/bin/env python3
"""
功能：输入参数时直接执行；
     未输入参数时，轮询读取 ADB 设备上 /data/local/tmp/query.txt 文件内容，
     并执行 python3 run_mobileagentv3.py --instruction "文件内容"

"""

import os
import uuid
import json
import time
import argparse
import subprocess
from PIL import Image
from datetime import datetime
from pathlib import Path
import signal
import sys

from utils.mobile_agent_e import (
    InfoPool, 
    Manager, 
    Executor, 
    Notetaker, 
    ActionReflector,
    INPUT_KNOW
)
import utils.controller as controller
from utils.call_mobile_agent_e import GUIOwlWrapper

myController = None

def run_instruction(adb_path, hdc_path, api_key, base_url, model, instruction, add_info, coor_type, if_notetaker, max_step=25, log_path="./logs", display=0):
    global myController
    if adb_path and hdc_path:
        raise ValueError("adb_path and hdc_path cannot be provided at the same time. Please specify only one of them.")
    if adb_path:
        from utils.android_controller import AndroidController
        myController = AndroidController(adb_path, display)
    else:
        from utils.harmonyos_controller import HarmonyOSController
        myController = HarmonyOSController(hdc_path)

    # 设置信号处理函数
    signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
    signal.signal(signal.SIGTERM, signal_handler) # 终止信号
    myController.show_border()

    if not os.path.exists(log_path):
        os.mkdir(log_path)
    now = datetime.now()
    time_str = now.strftime("%Y%m%d_%H%M%S")
    save_path = f"{log_path}/{time_str}_{instruction[:10]}"
    os.mkdir(save_path)
    image_save_path = os.path.join(save_path, "images")
    os.mkdir(image_save_path)
    info_pool = InfoPool(
        additional_knowledge_manager=add_info,
        additional_knowledge_executor=INPUT_KNOW,
        err_to_manager_thresh=2
    )
    vllm = GUIOwlWrapper(api_key, base_url, model)
    manager = Manager()
    executor = Executor()
    notetaker = Notetaker()
    action_reflector = ActionReflector()
    message_manager, message_operator, message_reflector, message_notekeeper = None, None, None, None
    info_pool.instruction = instruction
    print(f"max_step = {max_step}")
    for step in range(max_step):
        if step == max_step:
            task_result_path = os.path.join(save_path, "task_result.json")
            current_time = datetime.now()
            formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S.%f")
            task_result_data = {"goal": instruction, "finish_dtime": formatted_time, "hit_step_limit": 1.0}
            with open(task_result_path, 'w', encoding='utf-8') as json_file:
                json.dump(task_result_data, json_file, ensure_ascii=False, indent=4)
            break
        
        if step == 0:
            current_time = datetime.now()
            formatted_time = current_time.strftime(f'%Y-%m-%d-{current_time.hour * 3600 + current_time.minute * 60 + current_time.second}-{str(uuid.uuid4().hex[:8])}')
            local_image_dir = os.path.join(image_save_path, f"screenshot_{formatted_time}.png")
        else:
            local_image_dir = local_image_dir2
        # get the screenshot
        for _ in range(5):
            if not myController.get_screenshot(local_image_dir):
                print("Get screenshot failed, retry.")
                time.sleep(5)
            else:
                break
        
        width, height = Image.open(local_image_dir).size
        
        info_pool.error_flag_plan = False
        err_to_manager_thresh = info_pool.err_to_manager_thresh
        if len(info_pool.action_outcomes) >= err_to_manager_thresh:
            # check if the last err_to_manager_thresh actions are all errors
            latest_outcomes = info_pool.action_outcomes[-err_to_manager_thresh:]
            count = 0
            for outcome in latest_outcomes:
                if outcome in ["B", "C"]:
                    count += 1
            if count == err_to_manager_thresh:
                info_pool.error_flag_plan = True

        skip_manager = False
        ## if previous action is invalid, skip the manager and try again first ##
        if not info_pool.error_flag_plan and len(info_pool.action_history) > 0:
            if info_pool.action_history[-1]['action'] == 'invalid':
                skip_manager = True
        
        if not skip_manager:
            print("\n### Manager ... ###\n")
            prompt_planning = manager.get_prompt(info_pool)
            output_planning, message_manager, raw_response = vllm.predict_mm(
                prompt_planning,
                [local_image_dir]
            )
        
        message_save_path = os.path.join(save_path, f"step_{step+1}")
        os.mkdir(message_save_path)
        message_file = os.path.join(message_save_path, "manager.json")
        message_data = {"name": "manager", "messages": message_manager, "response": output_planning, "step_id": step+1}
        with open(message_file, 'w', encoding='utf-8') as json_file:
            json.dump(message_data, json_file, ensure_ascii=False, indent=4)

        parsed_result_planning = manager.parse_response(output_planning)
        info_pool.completed_plan = parsed_result_planning['completed_subgoal']
        info_pool.plan = parsed_result_planning['plan']
        if not raw_response:
            raise RuntimeError('Error calling vLLM in planning phase.')
        
        print('Completed subgoal: ' + info_pool.completed_plan)
        print('Planning thought: ' + parsed_result_planning['thought'])
        print('Plan: ' + info_pool.plan, "\n")
        
        if "Finished" in info_pool.plan.strip() and len(info_pool.plan.strip()) < 15:
            print("Instruction finished, stop the process.")
            task_result_path = os.path.join(save_path, "task_result.json")
            current_time = datetime.now()
            formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S.%f")
            task_result_data = {"goal": instruction, "finish_dtime": formatted_time, "hit_step_limit": 0.0}
            with open(task_result_path, 'w', encoding='utf-8') as json_file:
                json.dump(task_result_data, json_file, ensure_ascii=False, indent=4)
            break
        else:
            print("\n### Operator ... ###\n")

            prompt_action = executor.get_prompt(info_pool)
            output_action, message_operator, raw_response = vllm.predict_mm(
                prompt_action,
                [local_image_dir],
            )
            
            if not raw_response:
                raise RuntimeError('Error calling LLM in operator phase.')
            parsed_result_action = executor.parse_response(output_action)
            action_thought, action_object_str, action_description = parsed_result_action['thought'], parsed_result_action['action'], parsed_result_action['description']
            
            info_pool.last_action_thought = action_thought
            info_pool.last_summary = action_description
            
            if (not action_thought) or (not action_object_str):
                print('Action prompt output is not in the correct format.')
                info_pool.last_action = {"action": "invalid"}
                info_pool.action_history.append({"action": "invalid"})
                info_pool.summary_history.append(action_description)
                info_pool.action_outcomes.append("C")
                info_pool.error_descriptions.append("invalid action format, do nothing.")
                continue
        
        action_object_str = action_object_str.replace("```", "").replace("json", "").strip()
        print('Thought: ' + action_thought)
        print('Action: ' + action_object_str)
        print('Action description: ' + action_description)

        try:
            action_object = json.loads(action_object_str)
            operator_response = f'''### Thought ###
{action_thought}

### Action ###
{action_object}

### Description ###
{action_description}'''
            
            if action_object['action'] == "answer":
                message_file = os.path.join(message_save_path, "operator.json")
                message_data = {"name": "operator", "messages": message_operator, "response": operator_response, "step_id": step+1}
                with open(message_file, 'w', encoding='utf-8') as json_file:
                    json.dump(message_data, json_file, ensure_ascii=False, indent=4)

                answer_content = action_object['text']
                print(f"Instruction finished, answer: {answer_content}, stop the process.")
                task_result_path = os.path.join(save_path, "task_result.json")
                current_time = datetime.now()
                formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S.%f")
                task_result_data = {"goal": instruction, "finish_dtime": formatted_time, "hit_step_limit": 0.0}
                with open(task_result_path, 'w', encoding='utf-8') as json_file:
                    json.dump(task_result_data, json_file, ensure_ascii=False, indent=4)
                break
            
            if coor_type != "abs":
                if "coordinate" in action_object:
                    action_object['coordinate'] = [int(action_object['coordinate'][0] / 1000 * width), int(action_object['coordinate'][1] / 1000 * height)]
                if "coordinate2" in action_object:
                    action_object['coordinate2'] = [int(action_object['coordinate2'][0] / 1000 * width), int(action_object['coordinate2'][1] / 1000 * height)]
            
            if action_object['action'] == "click":
                myController.tap(action_object['coordinate'][0] * 2, action_object['coordinate'][1] * 2)
            elif action_object['action'] == "swipe":
                myController.slide(action_object['coordinate'][0], action_object['coordinate'][1], action_object['coordinate2'][0], action_object['coordinate2'][1])
            elif action_object['action'] == "type":
                myController.type(action_object['text'])
            elif action_object['action'] == "system_button":
                if action_object['button'] == "Back":
                    myController.back()
                elif action_object['button'] == "Home":
                    myController.home()
            
        except:
            info_pool.last_action = {"action": "invalid"}
            info_pool.action_history.append({"action": "invalid"})
            info_pool.summary_history.append(action_description)
            info_pool.action_outcomes.append("C")
            info_pool.error_descriptions.append("invalid action format, do nothing.")
            local_image_dir2 = local_image_dir
            continue
        
        message_file = os.path.join(message_save_path, "operator.json")
        message_data = {"name": "operator", "messages": message_operator, "response": operator_response, "step_id": step+1}
        with open(message_file, 'w', encoding='utf-8') as json_file:
            json.dump(message_data, json_file, ensure_ascii=False, indent=4)

        info_pool.last_action = json.loads(action_object_str)
        
        if step == 0:
            time.sleep(8) # maybe a pop-up when first open an app
        time.sleep(2)
        
        current_time = datetime.now()
        formatted_time = current_time.strftime(f'%Y-%m-%d-{current_time.hour * 3600 + current_time.minute * 60 + current_time.second}-{str(uuid.uuid4().hex[:8])}')
        local_image_dir2 = os.path.join(image_save_path, f"screenshot_{formatted_time}.png")
        
        # get the screenshot
        for _ in range(5):
            if not myController.get_screenshot(local_image_dir2):
                print("Get screenshot failed, retry.")
                time.sleep(5)
            else:
                break
        
        print("\n### Action Reflector ... ###\n")
        prompt_action_reflect = action_reflector.get_prompt(info_pool)
        output_action_reflect, message_reflector, raw_response = vllm.predict_mm(
            prompt_action_reflect,
            [
                local_image_dir,
                local_image_dir2,
            ],
        )
        
        message_file = os.path.join(message_save_path, "reflector.json")
        message_data = {"name": "reflector", "messages": message_reflector, "response": output_action_reflect, "step_id": step+1}
        with open(message_file, 'w', encoding='utf-8') as json_file:
            json.dump(message_data, json_file, ensure_ascii=False, indent=4)
        
        parsed_result_action_reflect = action_reflector.parse_response(output_action_reflect)
        outcome, error_description = (
            parsed_result_action_reflect['outcome'], 
            parsed_result_action_reflect['error_description']
        )
        progress_status = info_pool.completed_plan
        
        if "A" in outcome: # Successful. The result of the last action meets the expectation.
          action_outcome = "A"
        elif "B" in outcome: # Failed. The last action results in a wrong page. I need to return to the previous state.
            action_outcome = "B"
        elif "C" in outcome: # Failed. The last action produces no changes.
            action_outcome = "C"
        else:
            raise ValueError("Invalid outcome:", outcome)
        
        print('Action reflection outcome: ' + action_outcome)
        print('Action reflection error description: ' + error_description)
        print('Action reflection progress status: ' + progress_status, "\n")
        
        info_pool.action_history.append(json.loads(action_object_str))
        info_pool.summary_history.append(action_description)
        info_pool.action_outcomes.append(action_outcome)
        info_pool.error_descriptions.append(error_description)
        info_pool.progress_status = progress_status
        
        if action_outcome == "A" and if_notetaker:
            print("\n### NoteKeeper ... ###\n")
            prompt_note = notetaker.get_prompt(info_pool)
            output_note, message_notekeeper, raw_response = vllm.predict_mm(
                prompt_note,
                [local_image_dir2],
            )
            
            message_file = os.path.join(message_save_path, "notekeeper.json")
            message_data = {"name": "notekeeper", "messages": message_notekeeper, "response": output_note, "step_id": step+1}
            with open(message_file, 'w', encoding='utf-8') as json_file:
                json.dump(message_data, json_file, ensure_ascii=False, indent=4)
            
            parsed_result_note = notetaker.parse_response(output_note)
            important_notes = parsed_result_note['important_notes']
            info_pool.important_notes = important_notes

            print('Important notes: ' + important_notes, "\n")
    print("运行结束。")
    myController.hide_border()

def signal_handler(sig, frame):
    print('收到信号，正在退出...')
    myController.hide_border()
    time.sleep(1)
    sys.exit(0)

### adb轮询读取设备上文本内容 ###
class ADBFilePoller:
    def __init__(self, device_file_path="/data/local/tmp/query.txt", poll_interval=5):
        self.device_file_path = device_file_path
        self.poll_interval = poll_interval
        self.last_content = ""
        self.script_path = "run_mobileagentv3.py"

    def check_adb_devices(self):
        """检查是否有 ADB 设备连接"""
        try:
            result = subprocess.run(
                ["adb", "devices"],
                capture_output=True,
                text=True,
                check=True
            )
            lines = result.stdout.strip().split('\n')
            # 第一行是 "List of devices attached"，从第二行开始是设备信息
            devices = [line.split('\t')[0] for line in lines[1:] if line.strip()]
            return len(devices) > 0
        except (subprocess.CalledProcessError, FileNotFoundError):
            return False

    def read_file_from_device(self):
        """从设备读取文件内容"""
        try:
            result = subprocess.run(
                ["adb", "shell", f"cat {self.device_file_path}"],
                capture_output=True,
                text=True,
                encoding='utf-8',
                check=True
            )
            return result.stdout.strip()
        except subprocess.CalledProcessError as e:
            print(f"读取文件失败: {e}")
            return None

    def execute_mobile_agent(self, instruction):
        """执行 MobileAgent 脚本"""
        if not instruction or not instruction.strip():
            print("指令为空，跳过执行")
            return False

        try:
            # 检查脚本文件是否存在
            if not os.path.exists(self.script_path):
                print(f"错误: 脚本文件 {self.script_path} 不存在")
                return False

            print(f"执行指令: {instruction}")
            display = "0"
            if "副驾" in instruction:
                display = "1"
            command = [
                    "python",
                    self.script_path,
                    "--instruction",
                    instruction,
                    "--display",
                    display
                ]
#             run_instruction(args.adb_path, args.hdc_path, args.api_key, args.base_url, args.model, args.instruction, args.add_info, args.coor_type, args.notetaker, display = myDisplay)
#             run_instruction("adb", None, "sk-f29b7adf03ae4eeaa3565f98d5e1f650", "https://dashscope.aliyuncs.com/compatible-mode/v1", "qwen-vl-max",)
            print(f"执行命令：{' '.join(command)}")
            result = subprocess.run(
                command,
#                 capture_output=True,
#                 text=True,
                check=True
            )
            print(f"执行成功: {result.stdout}")
            return True
        except subprocess.CalledProcessError as e:
            print(f"执行失败: {e}")
            print(f"错误输出: {e.stderr}")
            return False

    def clear_file_on_device(self):
        """清空设备上的文件内容"""
        try:
            subprocess.run(
                ["adb", "shell", f"echo '' > {self.device_file_path}"],
                check=True
            )
            print("已清空设备文件内容")
        except subprocess.CalledProcessError as e:
            print(f"清空文件失败: {e}")

    def start_polling(self):
        """开始轮询"""
        print("启动 ADB")
        print(f"监控文件: {self.device_file_path}")
        print(f"间隔: {self.poll_interval} 秒")
        print("按 Ctrl+C 停止\n")

        if not self.check_adb_devices():
            print("错误: 未找到 ADB 设备，请确保:")
            print("1. USB 调试已开启")
            print("2. 设备已通过 USB 连接")
            print("3. 已授权此电脑的 ADB 调试")
            return

        print("检测到 ADB 设备，开始运行...")

        try:
            while True:
                content = self.read_file_from_device()
                print(f"检测到的内容: {content}")
                if content and content != self.last_content:
                    print(f"检测到新内容: {content}")

                    # 执行命令
                    success = self.execute_mobile_agent(content)

                    if success:
                        # 成功执行后更新最后内容并清空设备文件
                        self.last_content = content
                        self.clear_file_on_device()
                    else:
                        print("执行失败，保留文件内容")

                time.sleep(self.poll_interval)

        except KeyboardInterrupt:
            print("\n\n已停止")
        except Exception as e:
            print(f"发生错误: {e}")

def poller_main():
    # 检查 ADB 是否可用
    try:
        subprocess.run(["adb", "version"], capture_output=True, check=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("错误: ADB 未找到或不可用")
        print("请确保 ADB 已安装并在 PATH 环境变量中")
        sys.exit(1)

    # 创建轮询器实例
    poller = ADBFilePoller(
        device_file_path="/data/local/tmp/query.txt",
        poll_interval=3  # 3秒轮询一次
    )

    # 开始轮询
    poller.start_polling()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Run Mobile-Agent-v3 with a given model and instruction"
    )
    parser.add_argument("--adb_path", type=str, default="adb")
    parser.add_argument("--hdc_path", type=str)
    parser.add_argument("--api_key", type=str, default="sk-f29b7adf03ae4eeaa3565f98d5e1f650")
    parser.add_argument("--base_url", type=str, default="https://dashscope.aliyuncs.com/compatible-mode/v1")
    parser.add_argument("--model", type=str, default="qwen-vl-max")
    parser.add_argument("--instruction", type=str)
    parser.add_argument("--add_info", type=str, default="用中文描述思考和执行过程,应用列表控制在首页的右下角，回首页按钮在屏幕左下角，车鱼视听和抖音是同一个应用")
    parser.add_argument("--coor_type", type=str, default="abs")
    parser.add_argument("--notetaker", type=bool, default=False)
    parser.add_argument("--display", type=int, default=0)
    args = parser.parse_args()



    if args.instruction:
        myDisplay = args.display
        if "副驾" in args.instruction:
            myDisplay = 1
        run_instruction(args.adb_path, args.hdc_path, args.api_key, args.base_url, args.model, args.instruction, args.add_info, args.coor_type, args.notetaker, display = myDisplay)
    else:
        poller_main()
